//! Function literal (first-class function) code generation

use crate::ast::types::FunctionType;
use crate::ast::*;
use crate::error::{CompilerError, Result};
use inkwell::types::{BasicMetadataTypeEnum, BasicType, BasicTypeEnum, StructType};
use inkwell::values::{BasicValueEnum, FunctionValue, PointerValue};

use super::super::LLVMBackend;
use inkwell::AddressSpace;

struct CaptureField<'ctx> {
    name: String,
    nrc_type: Type,
    llvm_type: BasicTypeEnum<'ctx>,
    index: u32,
}

struct LambdaCaptureLayout<'ctx> {
    fields: Vec<CaptureField<'ctx>>,
    env_struct_type: Option<StructType<'ctx>>,
}

impl<'ctx> LLVMBackend<'ctx> {
    /// Generate LLVM IR for a function literal expression
    /// This allows functions to be treated as first-class values
    pub(crate) fn generate_function_literal_expression(
        &mut self,
        func_lit: &expr::FunctionLiteralExpr,
    ) -> Result<BasicValueEnum<'ctx>> {
        if let Some(function_name) = &func_lit.function_name {
            // This is a reference to a named function
            // Get the function pointer from the function map
            if let Some(func_value) = self.function_map.get(function_name) {
                // Return the function pointer as a value
                // In LLVM, function values are already pointers
                Ok(func_value.as_global_value().as_pointer_value().into())
            } else {
                Err(CompilerError::internal(&format!(
                    "Function '{}' not found",
                    function_name
                )))
            }
        } else if let Some(lambda) = &func_lit.lambda {
            let function_type = if let Some(Type::Function(func_type)) = &func_lit.function_type {
                func_type.clone()
            } else {
                return Err(CompilerError::internal(
                    "Lambda literals must include function type information",
                ));
            };

            let capture_layout = self.build_lambda_capture_layout(lambda)?;
            let capture_values = self.collect_capture_values(&capture_layout)?;
            let env_ptr = self.allocate_capture_environment(&capture_layout, &capture_values)?;

            let lambda_name = format!("lambda_{}", self.get_unique_lambda_id());
            let lambda_func =
                self.generate_lambda_function(&lambda_name, lambda, &capture_layout)?;

            self.build_closure_value(
                &function_type,
                lambda_func.as_global_value().as_pointer_value(),
                env_ptr,
            )
        } else {
            Err(CompilerError::internal(
                "FunctionLiteral must have either function_name or lambda",
            ))
        }
    }

    /// Generate a lambda function and return its FunctionValue
    fn generate_lambda_function(
        &mut self,
        lambda_name: &str,
        lambda: &expr::LambdaExpr,
        capture_layout: &LambdaCaptureLayout<'ctx>,
    ) -> Result<FunctionValue<'ctx>> {
        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        // Build parameter types
        let user_param_types: Vec<BasicTypeEnum> = lambda
            .parameters
            .iter()
            .map(|param| self.nrc_type_to_llvm_type(&Some(param.param_type.clone())))
            .collect::<Result<Vec<_>>>()?;

        // Convert to metadata types for function signature
        let mut metadata_param_types: Vec<BasicMetadataTypeEnum> =
            Vec::with_capacity(user_param_types.len() + 1);
        metadata_param_types.push(env_ptr_type.into());
        metadata_param_types.extend(
            user_param_types
                .iter()
                .map(|t| BasicMetadataTypeEnum::from(*t)),
        );

        // Build return type
        let return_type = if let Some(ref ret_type) = lambda.return_type {
            Some(self.nrc_type_to_llvm_type(&Some(ret_type.clone()))?)
        } else {
            None
        };

        // Create function type
        let fn_type = if let Some(ret_type) = return_type {
            ret_type.fn_type(&metadata_param_types, false)
        } else {
            self.context
                .void_type()
                .fn_type(&metadata_param_types, false)
        };

        // Create the function
        let function = self.module.add_function(lambda_name, fn_type, None);

        // Save current state
        let saved_function = self.current_function;
        let saved_return_type = self.current_function_return_type.take();
        let saved_insert_block = self.builder.get_insert_block();

        // Set up new function context
        self.current_function = Some(function);
        self.current_function_return_type = lambda.return_type.clone();

        // Create entry block
        let entry_block = self.context.append_basic_block(function, "entry");
        self.builder.position_at_end(entry_block);

        // Clear inherited variable state to avoid leaking outer variables
        let saved_variables = std::mem::take(&mut self.variable_map);
        let saved_variable_types = std::mem::take(&mut self.variable_type_map);

        // Add parameters to variable map
        for (i, param) in lambda.parameters.iter().enumerate() {
            let param_value = function.get_nth_param((i + 1) as u32).unwrap();
            param_value.set_name(&param.name);

            // Allocate stack space for the parameter
            let param_ptr = self
                .builder
                .build_alloca(user_param_types[i], &format!("{}_ptr", param.name))?;

            // Store the parameter value
            self.builder.build_store(param_ptr, param_value)?;

            // Add to variable map
            self.variable_map
                .insert(param.name.clone(), param_ptr.into());
            self.variable_type_map
                .insert(param.name.clone(), param.param_type.clone());
        }

        // Materialize captured variables from the environment
        if let Some(env_struct_type) = capture_layout.env_struct_type {
            let env_param = function.get_nth_param(0).unwrap();
            env_param.set_name("__env");
            let env_param_ptr = env_param.into_pointer_value();
            let typed_env_ptr = self
                .builder
                .build_bitcast(
                    env_param_ptr,
                    env_struct_type.ptr_type(AddressSpace::default()),
                    "env_struct_ptr",
                )?
                .into_pointer_value();

            for field in &capture_layout.fields {
                let field_ptr = unsafe {
                    self.builder.build_gep(
                        env_struct_type,
                        typed_env_ptr,
                        &[
                            self.context.i32_type().const_int(0, false),
                            self.context.i32_type().const_int(field.index as u64, false),
                        ],
                        &format!("{}_capture_ptr", field.name),
                    )?
                };
                let loaded = self
                    .builder
                    .build_load(field.llvm_type, field_ptr, &field.name)?;
                let local_alloca = self
                    .builder
                    .build_alloca(field.llvm_type, &format!("{}_capture", field.name))?;
                self.builder.build_store(local_alloca, loaded)?;
                self.variable_map
                    .insert(field.name.clone(), local_alloca.into());
                self.variable_type_map
                    .insert(field.name.clone(), field.nrc_type.clone());
            }
        }

        // Generate function body
        let body_value = self.generate_expression(&lambda.body)?;

        // Build return only if the current block doesn't already have a terminator
        // (e.g., if the lambda body contains an explicit return statement)
        let current_block = self.builder.get_insert_block();
        let needs_return = if let Some(block) = current_block {
            block.get_terminator().is_none()
        } else {
            true
        };

        if needs_return {
            if lambda.return_type.is_some() {
                self.builder.build_return(Some(&body_value))?;
            } else {
                self.builder.build_return(None)?;
            }
        }

        // IMPORTANT: Restore builder position BEFORE restoring current_function
        // to avoid any chance of adding instructions to the lambda function
        if let Some(block) = saved_insert_block {
            self.builder.position_at_end(block);
        } else {
            // If there's no saved block, clear the builder position
            self.builder.clear_insertion_position();
        }

        // Now restore previous state
        self.current_function = saved_function;
        self.current_function_return_type = saved_return_type;
        self.variable_map = saved_variables;
        self.variable_type_map = saved_variable_types;

        // Add to function map so it can be called later
        self.function_map.insert(lambda_name.to_string(), function);

        Ok(function)
    }

    /// Get a unique ID for lambda naming
    fn get_unique_lambda_id(&self) -> usize {
        // Count how many lambdas we've generated so far
        self.function_map
            .keys()
            .filter(|k| k.starts_with("lambda_"))
            .count()
    }

    /// Generate code for calling a function pointer
    /// This is used when a function value is called
    pub(crate) fn generate_function_pointer_call(
        &mut self,
        func_ptr: BasicValueEnum<'ctx>,
        arguments: &[Expression],
        func_type_info: Option<&FunctionType>,
        _location: crate::error::Location,
    ) -> Result<BasicValueEnum<'ctx>> {
        let arg_values: Vec<BasicValueEnum> = arguments
            .iter()
            .map(|arg| self.generate_expression(arg))
            .collect::<Result<Vec<_>>>()?;

        let closure_value = match func_ptr {
            BasicValueEnum::StructValue(struct_val) => struct_val,
            BasicValueEnum::PointerValue(ptr_val) => {
                let func_type = func_type_info.ok_or_else(|| {
                    CompilerError::internal(
                        "Missing function type information for closure pointer value",
                    )
                })?;
                let closure_type =
                    self.nrc_type_to_llvm_type(&Some(Type::Function(func_type.clone())))?;
                let struct_type = closure_type.into_struct_type();
                self.builder
                    .build_load(struct_type, ptr_val, "closure_load")?
                    .into_struct_value()
            }
            _ => {
                return Err(CompilerError::internal(
                    "Function values must be closure structs",
                ))
            }
        };

        let fn_ptr_value = self
            .builder
            .build_extract_value(closure_value, 0, "closure_fn")?
            .into_pointer_value();
        let env_ptr_value = self
            .builder
            .build_extract_value(closure_value, 1, "closure_env")?
            .into_pointer_value();

        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());

        let (fn_ptr_type, param_types, return_type, returns_void) = if let Some(func_type) =
            func_type_info
        {
            let fn_ptr_type = self.closure_code_pointer_type(func_type)?;
            let mut params = Vec::with_capacity(func_type.parameter_types.len() + 1);
            params.push(env_ptr_type.into());
            for param in &func_type.parameter_types {
                params.push(self.nrc_type_to_llvm_type(&Some(param.clone()))?.into());
            }
            let (ret_type, is_void) = if let Some(ret) = &func_type.return_type {
                (
                    Some(self.nrc_type_to_llvm_type(&Some(*ret.clone()))?),
                    false,
                )
            } else {
                (None, true)
            };
            (fn_ptr_type, params, ret_type, is_void)
        } else {
            let mut params: Vec<BasicMetadataTypeEnum> = Vec::with_capacity(arg_values.len() + 1);
            params.push(env_ptr_type.into());
            params.extend(
                arg_values
                    .iter()
                    .map(|v| BasicMetadataTypeEnum::from(v.get_type())),
            );
            let ret_type = self.context.i32_type();
            (
                ret_type
                    .fn_type(&params, false)
                    .ptr_type(AddressSpace::default()),
                params,
                Some(ret_type.into()),
                false,
            )
        };

        let fn_ptr_cast = self
            .builder
            .build_bitcast(fn_ptr_value, fn_ptr_type, "closure_fn_cast")?
            .into_pointer_value();

        let env_arg = self
            .builder
            .build_bitcast(env_ptr_value, env_ptr_type, "closure_env_cast")?
            .into_pointer_value();

        let mut call_values: Vec<BasicValueEnum> = Vec::new();
        call_values.push(env_arg.into());
        for value in &arg_values {
            call_values.push(*value);
        }

        let call_metadata: Vec<_> = call_values.iter().map(|v| (*v).into()).collect();

        let fn_type = if let Some(ret) = return_type {
            ret.fn_type(&param_types, false)
        } else {
            self.context.void_type().fn_type(&param_types, false)
        };

        let call_result = self.builder.build_indirect_call(
            fn_type,
            fn_ptr_cast,
            &call_metadata,
            "closure_call",
        )?;

        if returns_void {
            Ok(self.context.i64_type().const_int(0, false).into())
        } else if let Some(result) = call_result.try_as_basic_value().left() {
            Ok(result)
        } else {
            Err(CompilerError::internal(
                "Expected return value from closure call",
            ))
        }
    }

    pub(super) fn build_function_value_reference(
        &mut self,
        name: &str,
    ) -> Result<BasicValueEnum<'ctx>> {
        if self.overloaded_function_names.contains(name) {
            return Err(CompilerError::internal(format!(
                "Cannot take overloaded function '{}' as a value yet",
                name
            )));
        }

        let func_decl = self.function_decls.get(name).cloned().ok_or_else(|| {
            CompilerError::internal(format!(
                "Function '{}' not found for first-class reference",
                name
            ))
        })?;

        let symbol_name = self.function_symbol_name(&func_decl);
        let function_value = *self.function_map.get(&symbol_name).ok_or_else(|| {
            CompilerError::internal(format!(
                "Function '{}' missing from LLVM module",
                symbol_name
            ))
        })?;

        let shim = self.get_or_create_closure_shim(&func_decl, function_value)?;

        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let null_env = env_ptr_type.const_null();
        let function_type = Self::function_decl_to_type(&func_decl);
        self.build_closure_value(
            &function_type,
            shim.as_global_value().as_pointer_value(),
            null_env,
        )
    }

    fn build_lambda_capture_layout(
        &mut self,
        lambda: &expr::LambdaExpr,
    ) -> Result<LambdaCaptureLayout<'ctx>> {
        let mut fields = Vec::new();
        for (idx, capture) in lambda.captures.iter().enumerate() {
            let resolved_type = if let Some(var_type) = self.variable_type_map.get(&capture.name) {
                var_type.clone()
            } else {
                capture.var_type.clone()
            };
            let llvm_type = self.nrc_type_to_llvm_type(&Some(resolved_type.clone()))?;
            fields.push(CaptureField {
                name: capture.name.clone(),
                nrc_type: resolved_type,
                llvm_type,
                index: idx as u32,
            });
        }

        let env_struct_type = if fields.is_empty() {
            None
        } else {
            let llvm_fields: Vec<BasicTypeEnum> = fields.iter().map(|f| f.llvm_type).collect();
            Some(self.context.struct_type(&llvm_fields, false))
        };

        Ok(LambdaCaptureLayout {
            fields,
            env_struct_type,
        })
    }

    fn collect_capture_values(
        &mut self,
        layout: &LambdaCaptureLayout<'ctx>,
    ) -> Result<Vec<BasicValueEnum<'ctx>>> {
        let mut values = Vec::new();
        for field in &layout.fields {
            let base_value = if let Some(var_value) = self.variable_map.get(&field.name) {
                if var_value.is_pointer_value() {
                    let ptr = var_value.into_pointer_value();
                    self.builder.build_load(
                        field.llvm_type,
                        ptr,
                        &format!("{}_capture_load", field.name),
                    )?
                } else {
                    *var_value
                }
            } else if let Some(global) = self.module.get_global(&field.name) {
                let global_ptr = global.as_pointer_value();
                self.builder.build_load(
                    field.llvm_type,
                    global_ptr,
                    &format!("{}_global_capture", field.name),
                )?
            } else {
                return Err(CompilerError::internal(format!(
                    "Cannot capture undefined variable '{}'",
                    field.name
                )));
            };

            let cloned = self.clone_capture_value(base_value, &field.nrc_type)?;
            values.push(cloned);
        }
        Ok(values)
    }

    fn clone_capture_value(
        &mut self,
        value: BasicValueEnum<'ctx>,
        ty: &Type,
    ) -> Result<BasicValueEnum<'ctx>> {
        match ty {
            Type::Vec(vec_type) => self.deep_copy_slice(value, vec_type),
            _ => Ok(value),
        }
    }

    fn allocate_capture_environment(
        &mut self,
        layout: &LambdaCaptureLayout<'ctx>,
        values: &[BasicValueEnum<'ctx>],
    ) -> Result<PointerValue<'ctx>> {
        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());

        if layout.fields.is_empty() {
            return Ok(env_ptr_type.const_null());
        }

        let env_struct_type = layout.env_struct_type.ok_or_else(|| {
            CompilerError::internal("Missing environment struct type for captured values")
        })?;

        let struct_size = self.get_llvm_type_size(env_struct_type.into());
        let size_const = self.context.i64_type().const_int(struct_size as u64, false);
        let size_value: BasicValueEnum = size_const.into();
        let raw_ptr = match self.generate_malloc(&size_value)? {
            BasicValueEnum::PointerValue(ptr) => ptr,
            _ => {
                return Err(CompilerError::internal(
                    "malloc for capture environment did not return a pointer",
                ))
            }
        };

        let typed_env_ptr = self
            .builder
            .build_bitcast(
                raw_ptr,
                env_struct_type.ptr_type(AddressSpace::default()),
                "env_struct_alloc",
            )?
            .into_pointer_value();

        for (idx, field) in layout.fields.iter().enumerate() {
            let field_ptr = unsafe {
                self.builder.build_gep(
                    env_struct_type,
                    typed_env_ptr,
                    &[
                        self.context.i32_type().const_int(0, false),
                        self.context.i32_type().const_int(idx as u64, false),
                    ],
                    &format!("env_field_{}_ptr", field.name),
                )?
            };
            self.builder.build_store(field_ptr, values[idx])?;
        }

        let env_i8_ptr = self
            .builder
            .build_bitcast(typed_env_ptr, env_ptr_type, "env_i8_ptr")?
            .into_pointer_value();
        Ok(env_i8_ptr)
    }

    fn build_closure_value(
        &mut self,
        func_type: &FunctionType,
        function_ptr: PointerValue<'ctx>,
        env_ptr: PointerValue<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        let closure_type = self.nrc_type_to_llvm_type(&Some(Type::Function(func_type.clone())))?;
        let struct_type = closure_type.into_struct_type();
        let fn_ptr_type = self.closure_code_pointer_type(func_type)?;
        let fn_ptr_cast = self
            .builder
            .build_bitcast(function_ptr, fn_ptr_type, "closure_fn_store")?
            .into_pointer_value();
        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let env_cast = self
            .builder
            .build_bitcast(env_ptr, env_ptr_type, "closure_env_store")?
            .into_pointer_value();

        let mut closure_value = struct_type.get_undef();
        closure_value = self
            .builder
            .build_insert_value(closure_value, fn_ptr_cast, 0, "set_fn")?
            .into_struct_value();
        closure_value = self
            .builder
            .build_insert_value(closure_value, env_cast, 1, "set_env")?
            .into_struct_value();
        Ok(closure_value.into())
    }

    fn get_or_create_closure_shim(
        &mut self,
        func_decl: &FunctionDecl,
        target: FunctionValue<'ctx>,
    ) -> Result<FunctionValue<'ctx>> {
        let symbol_name = self.function_symbol_name(func_decl);
        if let Some(existing) = self.function_closure_shims.get(&symbol_name) {
            return Ok(*existing);
        }

        let env_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());
        let mut param_types: Vec<BasicMetadataTypeEnum> =
            Vec::with_capacity(func_decl.parameters.len() + 1);
        param_types.push(env_ptr_type.into());
        for param in &func_decl.parameters {
            param_types.push(
                self.nrc_type_to_llvm_type(&Some(param.param_type.clone()))?
                    .into(),
            );
        }

        let shim_type = if let Some(ret_type) = &func_decl.return_type {
            let ret = self.nrc_type_to_llvm_type(&Some(ret_type.clone()))?;
            ret.fn_type(&param_types, false)
        } else {
            self.context.void_type().fn_type(&param_types, false)
        };

        let shim_name = format!("closure.shim.{}", symbol_name);
        let shim_function = self.module.add_function(&shim_name, shim_type, None);

        let saved_block = self.builder.get_insert_block();
        let saved_function = self.current_function;
        self.current_function = Some(shim_function);

        let entry = self.context.append_basic_block(shim_function, "entry");
        self.builder.position_at_end(entry);

        let mut call_args: Vec<BasicValueEnum> = Vec::new();
        for i in 0..func_decl.parameters.len() {
            let param_value = shim_function.get_nth_param((i + 1) as u32).unwrap();
            call_args.push(param_value.into());
        }

        let call_metadata: Vec<_> = call_args.iter().map(|v| (*v).into()).collect();
        let call_site = self
            .builder
            .build_call(target, &call_metadata, "shim_call")?;

        if func_decl.return_type.is_some() {
            let ret_val = call_site
                .try_as_basic_value()
                .left()
                .ok_or_else(|| CompilerError::internal("Expected return value in shim"))?;
            self.builder.build_return(Some(&ret_val))?;
        } else {
            self.builder.build_return(None)?;
        }

        if let Some(block) = saved_block {
            self.builder.position_at_end(block);
        } else {
            self.builder.clear_insertion_position();
        }
        self.current_function = saved_function;

        self.function_closure_shims
            .insert(symbol_name.clone(), shim_function);
        Ok(shim_function)
    }

    fn function_decl_to_type(func_decl: &FunctionDecl) -> FunctionType {
        FunctionType {
            parameter_types: func_decl
                .parameters
                .iter()
                .map(|p| p.param_type.clone())
                .collect(),
            return_type: func_decl
                .return_type
                .as_ref()
                .map(|ty| Box::new(ty.clone())),
            variadic: false,
            location: func_decl.location,
        }
    }
}
